package com.huxiaohui.ssbx.websocket;

import cn.hutool.core.thread.ThreadUtil;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.huxiaohui.ssbx.agent.Agent;
import com.huxiaohui.ssbx.agent.AgentFactory;
import com.huxiaohui.ssbx.agent.enums.AgentTypeEnum;
import com.huxiaohui.ssbx.constants.ResponseConstants;
import com.huxiaohui.ssbx.constants.TtsConstants;
import com.huxiaohui.ssbx.service.IVoiceService;
import com.huxiaohui.ssbx.service.impl.SessionServiceImpl;
import com.huxiaohui.ssbx.utils.ResponseFutureHolder;
import com.huxiaohui.ssbx.utils.SessionHolder;
import com.huxiaohui.ssbx.utils.SessionMessage;
import jakarta.annotation.Resource;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;

import java.io.IOException;
import java.io.ByteArrayOutputStream;
import java.nio.ByteBuffer;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

/**
 * 语音WebSocket处理器
 */
@Component
@RequiredArgsConstructor
@Slf4j
public class WebSocketHandler implements org.springframework.web.socket.WebSocketHandler {

    private final IVoiceService voiceService;
    private final SessionServiceImpl sessionService;
    private final ExecutorService executor = Executors.newCachedThreadPool();
    @Resource(name = "sendMessageExecutor")
    private ThreadPoolTaskExecutor sendMessageExecutor;


    // 聚合音频分片缓冲区（按会话）
    private final Map<String, ByteArrayOutputStream> audioBuffers = new ConcurrentHashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String sessionId = getSessionIdFromUri(session.getUri().toString());
        if (sessionId != null && sessionService.isValidSession(sessionId)) {
            SessionHolder.put(sessionId, session);
            log.info("WebSocket连接建立，SessionID: {}", sessionId);

            // 发送连接成功消息
            SessionMessage.sendMessage(session, "连接成功", "system");
        } else {
            log.warn("无效或过期的SessionID，关闭连接: {}", sessionId);
            SessionMessage.sendMessage(session, "会话过期，请重启应用程序", "error");
            session.close();
        }
    }

    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
        String sessionId = getSessionIdFromUri(session.getUri().toString());
        if (sessionId == null) {
            return;
        }

        if (message instanceof TextMessage) {
            // 处理文本消息
            String text = ((TextMessage) message).getPayload();
            log.debug("收到文本消息: {}", text);
            handleTextMessage(sessionId, text);
        } else if (message instanceof BinaryMessage binMsg) {
            // 处理音频数据（支持分片聚合）
            ByteBuffer payload = binMsg.getPayload();
            byte[] chunk = new byte[payload.remaining()];
            payload.get(chunk);
            log.debug("收到音频数据分片，大小: {} bytes, isLast: {}", chunk.length, binMsg.isLast());

            ByteArrayOutputStream buffer = audioBuffers.computeIfAbsent(sessionId, k -> new ByteArrayOutputStream());
            try {
                buffer.write(chunk);
            } catch (IOException e) {
                log.error("写入音频分片失败", e);
            }

            if (binMsg.isLast()) {
                byte[] completeAudio = buffer.toByteArray();
                audioBuffers.remove(sessionId);
                log.debug("完整音频聚合完成，大小: {} bytes", completeAudio.length);
                handleAudioMessage(sessionId, completeAudio);
            }
        }
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        log.error("WebSocket传输错误", exception);
        // 清理会话
        String sessionId = getSessionIdFromUri(session.getUri().toString());
        if (sessionId != null) {
            SessionHolder.remove(sessionId);
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
        String sessionId = getSessionIdFromUri(session.getUri().toString());
        if (sessionId != null) {
            SessionHolder.remove(sessionId);
            log.info("WebSocket连接关闭，SessionID: {}, 关闭状态: {}", sessionId, closeStatus);
        }
    }

    @Override
    public boolean supportsPartialMessages() {
        return true;
    }

    /**
     * 从URI中提取SessionID
     */
    private String getSessionIdFromUri(String uri) {
        try {
            String[] parts = uri.split("sessionId=");
            if (parts.length > 1) {
                return parts[1];
            }
        } catch (Exception e) {
            log.error("解析SessionID失败", e);
        }
        return null;
    }

    /**
     * 处理文本消息
     */
    private void handleTextMessage(String sessionId, String text) {
        try {
            // 处理文本
            log.info("处理文本消息，SessionID: {}, 内容: {}", sessionId, text);

            JSONObject jsonObject = JSONUtil.parseObj(text);
            String requestId = jsonObject.getStr("requestId");
            String content = jsonObject.getStr("content");

            CompletableFuture<String> future = ResponseFutureHolder.get(requestId);
            if (future != null && !future.isDone()) {
                future.complete(content);
            }
        } catch (Exception e) {
            log.error("处理文本消息失败", e);
            SessionMessage.sendMessage(SessionHolder.get(sessionId), "处理失败，请重试", "error");
        }
    }

    /**
     * 处理音频消息
     */
    private void handleAudioMessage(String sessionId, byte[] audioData) {
        try {
            // 检查音频数据大小，避免处理过小的数据
            if (audioData.length < 111111) {
                log.debug("音频数据太小，跳过处理: {} bytes", audioData.length);
                SessionMessage.sendMessage(SessionHolder.get(sessionId), "时长太短", "tips");
                return;
            }

            // 响应用户语音
            sendMessageExecutor.execute(() -> {
                SessionMessage.sendAudioMessage(SessionHolder.get(sessionId), voiceService.generateTTS(sessionId, TtsConstants.FIRST_VOICE));
            });
            // 调用语音服务处理音频-->转文字
            String text = voiceService.processAudio(sessionId, audioData);
            if (text.contains(ResponseConstants.ERROR_TAG)) {
                // 响应错误消息
                SessionMessage.sendMessage(SessionHolder.get(sessionId), ResponseConstants.msg(text), "error");
            } else {
                // 处理任务
                Agent agent = AgentFactory.getAgent(AgentTypeEnum.RE_ACT_PLANNING_AGENT).apply(sessionId);
                Runnable task = () -> {
                    // 处理任务
                    String result = agent.solveTask(text);
                    log.info("处理任务结果: {}", result);
                    SessionMessage.sendAudioMessage(SessionHolder.get(sessionId), voiceService.generateTTS(sessionId, result));
                    // 发送语音---> 任务结束语音
//                    ThreadUtil.sleep(3000);
//                    byte[] ttsAudio = voiceService.generateTTS();
//                    SessionMessage.sendAudioMessage(SessionHolder.get(sessionId), ttsAudio);
//                    log.info("任务结束，发送结束语音");
                };
                executor.execute(task);
            }


        } catch (Exception e) {
            log.error("处理音频消息失败", e);
            SessionMessage.sendMessage(SessionHolder.get(sessionId), "处理失败", "error");
            // 发送语音---> 任务结束语音
            byte[] ttsAudio = voiceService.generateTTS();
            SessionMessage.sendAudioMessage(SessionHolder.get(sessionId), ttsAudio);
            log.info("任务结束，发送结束语音");
        }
    }

//    /**
//     * 发送文本消息
//     */
//    public void sendMessage(WebSocketSession session, String message, String type) {
//        if (session != null && session.isOpen()) {
//            try {
//                String jsonMessage = objectMapper.writeValueAsString(Map.of(
//                    "type", type,
//                    "message", message,
//                    "timestamp", System.currentTimeMillis()
//                ));
//                session.sendMessage(new TextMessage(jsonMessage));
//            } catch (IOException e) {
//                log.error("发送消息失败", e);
//            }
//        }
//    }
//
//    /**
//     * 发送音频消息
//     */
//    private void sendAudioMessage(WebSocketSession session, byte[] audioData) {
//        if (session != null && session.isOpen()) {
//            try {
//                session.sendMessage(new BinaryMessage(audioData));
//            } catch (IOException e) {
//                log.error("发送音频消息失败", e);
//            }
//        }
//    }
}
